"""Paint a small landscape using joint targets, palette contacts, and RGB checks.""" import argparse import json from pathlib import Path import time import urllib.request import cv2 import numpy as np from PIL import Image from arm_kinematics import SO101Kinematics ROOT = Path(__file__).resolve().parents[1] class Painter: def __init__(self, port, results_dir=None): self.url = f'http://127.0.0.1:{port}' self.config = json.loads((ROOT/'config.json').read_text()) self.kin = SO101Kinematics(ROOT/'scene/scene.xml') self.grip = self.config['gripper_angle'] self.hover = self.config['hover_height'] self.z = self.config['contact_height'] self.H = None self.history = [] self.results = Path(results_dir) if results_dir else ROOT/'results' self.results.mkdir(parents=True, exist_ok=True) def state(self): with urllib.request.urlopen(self.url+'/state', timeout=10) as r: return json.load(r) def command(self, endpoint, **payload): req = urllib.request.Request(self.url+endpoint, data=json.dumps(payload).encode(), headers={'Content-Type':'application/json'}) with urllib.request.urlopen(req, timeout=180) as r: result = json.load(r) if 'error' in result: raise RuntimeError(result['error']) if any(result.get('warning_counts', [])): raise RuntimeError('Physics warning: '+str(result['warning_counts'])) return result def joints(self, q, seconds=1., label='move'): return self.command('/move', target=np.asarray(q).tolist(), seconds=seconds, label=label) def path(self, xyz, seconds, label): seed = self.state()['qpos'] targets = [] for point in xyz: seed = self.kin.solve(point, seed=seed, grip=self.grip) targets.append(seed.tolist()) return self.command('/trajectory', targets=targets, seconds=seconds, label=label) def line(self, target, seconds=1., label='move'): start = self.kin.forward_grasp(self.state()['qpos']) count = max(2, int(np.ceil(np.linalg.norm(np.asarray(target)-start)/.005))) return self.path(np.linspace(start, target, count+1)[1:], seconds, label) def observe(self, label): result = self.command('/camera', label=label) rgb = np.asarray(Image.open(result['image']).convert('RGB')) return rgb, result def localize_canvas(self): rgb, result = self.observe('locate-canvas') hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV) mask = ((hsv[:,:,0]>=85)&(hsv[:,:,0]<=99)&(hsv[:,:,1]>130)&(hsv[:,:,2]>120)).astype(np.uint8) count, _, stats, centers = cv2.connectedComponentsWithStats(mask, connectivity=8) valid = [i for i in range(1,count) if 5 < stats[i,cv2.CC_STAT_AREA] < 1500] valid = sorted(valid,key=lambda i:stats[i,cv2.CC_STAT_AREA],reverse=True)[:4] if len(valid)!=4: raise RuntimeError(f'Expected four visible cyan canvas markers, found {len(valid)}') points = centers[valid].astype(np.float32) indexes = [np.argmin(points[:,0]),np.argmax(points[:,1]),np.argmax(points[:,0]),np.argmin(points[:,1])] if len(set(indexes))!=4: raise RuntimeError('Canvas marker ordering is ambiguous') image_points = points[indexes] world = np.array(self.config['fiducials'],np.float32)[:,:2] H, _ = cv2.findHomography(image_points,world) b=self.config['canvas_bounds'] W = np.array([[511/(b['xmax']-b['xmin']),0,-511*b['xmin']/(b['xmax']-b['xmin'])], [0,-511/(b['ymax']-b['ymin']),511*b['ymax']/(b['ymax']-b['ymin'])],[0,0,1]]) self.H = W @ H report=dict(method='Four cyan marker centroids in the front RGB image',image_points=image_points.tolist(), known_marker_xy=world.tolist(),camera_to_canvas=self.H.tolist()) (self.results/'camera-localization.json').write_text(json.dumps(report,indent=2)) print('Located the canvas from all four RGB markers',flush=True) def inspect(self,label): rgb,state=self.observe(label) canvas=cv2.warpPerspective(rgb,self.H,(512,512)) hsv=cv2.cvtColor(canvas,cv2.COLOR_RGB2HSV) h,s,v=cv2.split(hsv) vivid=(s>45)&(v>35) inside=np.zeros(h.shape,bool);inside[22:-22,22:-22]=True masks={'blue':(h>=95)&(h<=125),'green':(h>=40)&(h<=85),'orange':(h>=7)&(h<=23),'red':(h<=6)|(h>=172)} pixels={name:int(np.count_nonzero(vivid&selected&inside)) for name,selected in masks.items()} Image.fromarray(canvas).save(self.results/'camera-canvas.png') self.history.append(dict(label=label,pixels=pixels,contact_frames=state['contact_frames'],paint=state['paint'])) print(f'Camera check {label}: colored pixels {pixels}',flush=True) return pixels,state def xy(self,u,v): b=self.config['canvas_bounds'] return [b['xmin']+u*(b['xmax']-b['xmin']), b['ymax']-v*(b['ymax']-b['ymin'])] def load(self,color): x,y,z=self.config['palette'][color]['position'] q=self.state()['qpos'] current=self.kin.forward_grasp(q) self.line([*current[:2],self.hover],.6,'lift-brush') self.line([x,y,self.hover],1.4,'approach-'+color+'-paint') self.line([x,y,z-.002],.8,'dip-'+color) self.joints(self.state()['ctrl'],.35,'load-'+color) state=self.state() if state['brush_color']!=color or state['paint']['load'] < .9: raise RuntimeError(f'No physical {color} palette contact: {state}') self.line([x,y,self.hover],.65,'lift-loaded-brush') print(f'Loaded {color} from its paint well',flush=True) def stroke(self,points,label): xy=np.asarray([self.xy(u,v) for u,v in points]) self.line([*xy[0],self.hover],1.15,'approach-'+label) self.line([*xy[0],self.z],.65,'touch-'+label) before=self.state()['contact_frames'] length=float(np.linalg.norm(np.diff(xy,axis=0),axis=1).sum()) seconds=max(.8,length/.035) # Densify the path in Cartesian space before solving joint targets. trajectory=[] for start,end in zip(xy,xy[1:]): n=max(2,int(np.ceil(np.linalg.norm(end-start)/.003))) trajectory.extend([[*p,self.z] for p in np.linspace(start,end,n+1)[1:]]) result=self.path(trajectory,seconds,label) if result['contact_frames'] <= before: raise RuntimeError(f'Brush missed the paper on {label}') self.line([*xy[-1],self.hover],.65,'lift-'+label) return self.inspect(label) def main(): parser=argparse.ArgumentParser() parser.add_argument('--port',type=int,default=8879) parser.add_argument('--record',action='store_true') parser.add_argument('--probe',action='store_true',help='Paint one blue test stroke only.') args=parser.parse_args() p=Painter(args.port) (ROOT/'results').mkdir(exist_ok=True) # Start with a vertical brush away from the cyan corner markers. initial=p.kin.solve([.18,-.12,p.hover],seed=p.state()['qpos'],grip=p.grip) p.joints(initial,1.2,'prepare-palette') p.localize_canvas() recording=False try: if args.record: p.command('/record/start');recording=True p.load('blue') if args.probe: p.stroke([[.20,.30],[.75,.30]],'blue-test-stroke') else: # Planned landscape geometry; all visible marks come from physical # brush contact. RGB observations check progress after each stroke. for i,v in enumerate((.20,.245,.29,.335)): points=[(.13+.75*t, v+.014*np.sin(t*np.pi*2)) for t in np.linspace(0,1,28)] p.stroke(points,f'sky-{i+1}') p.load('blue') for i,v in enumerate((.76,.81,.86)): p.stroke([(t,v+.018*np.sin(t*9)) for t in np.linspace(.14,.86,28)],f'lake-{i+1}') p.load('green') p.stroke([[.12,.65],[.35,.40],[.59,.65]],'left-mountain') p.stroke([[.36,.66],[.64,.42],[.89,.65]],'right-mountain') for i,(left,peak,right) in enumerate(((.18,.46,.53),(.24,.52,.48))): p.stroke([[left,.65],[.35,peak],[right,.65]],f'left-mountain-wash-{i+1}') for i,(left,peak,right) in enumerate(((.43,.48,.83),(.49,.54,.78))): p.stroke([[left,.66],[.64,peak],[right,.65]],f'right-mountain-wash-{i+1}') p.stroke([[.13,.66],[.48,.64],[.87,.67]],'shore') p.load('orange') center=np.array([.72,.17]) # A spiral fills the small sun with actual overlapping brushstrokes. sun=[(center+np.array([np.cos(t),np.sin(t)])*r).tolist() for t,r in zip(np.linspace(0,np.pi*6,100),np.linspace(.053,.006,100))] p.stroke(sun,'sun') p.line([.18,-.12,p.hover],1.4,'park-brush') pixels,state=p.inspect('finished') expected=['blue'] if args.probe else ['blue','green','orange'] if any(pixels[color]<80 for color in expected): raise RuntimeError(f'Final front-camera check is missing paint colors: {pixels}') report=dict(result='passed',picture='watercolor mountain landscape',camera_color_pixels=pixels, observations=p.history,final_state=state, method='Preplanned strokes with RGB progress checks. Paint transfer requires MuJoCo brush contact; no fluid solver.') (ROOT/'results/run-report.json').write_text(json.dumps(report,indent=2)) print('Finished painting; all expected colors visible in the camera.',flush=True) if recording: time.sleep(2.) finally: if recording: print(p.command('/record/stop'),flush=True) if __name__=='__main__': main()