clip-image

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoint/sam_vit_h_4b8939.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,121 @@
+import gradio as gr
+import numpy as np
+from meta_segment_anything import SegmentAnything
+from PIL import Image, ImageDraw
+
+def check_location(image, enable1, left1, top1, enable2, left2, top2, enable3, left3, top3):
+	if image is None:
+		yield None
+		return
+	if not enable1 and not enable2 and not enable3:
+		yield None
+		return
+
+	points = []
+	if enable1:
+		points.append([left1, top1])
+	if enable2:
+		points.append([left2, top2])
+	if enable3:
+		points.append([left3, top3])
+	for point in points:
+		left, top = point
+		draw = ImageDraw.Draw(image)
+		draw.ellipse([(left - 2, top - 2), (left + 3, top + 3)], fill=(255, 0, 0))
+		yield image
+
+def process_image(image, enable1, left1, top1, enable2, left2, top2, enable3, left3, top3):
+	if image is None:
+		yield None
+		return
+	if not enable1 and not enable2 and not enable3:
+		yield None
+		return
+
+	predictor = SegmentAnything()
+	points = []
+	if enable1:
+		points.append([left1, top1])
+	if enable2:
+		points.append([left2, top2])
+	if enable3:
+		points.append([left3, top3])
+	newImage = Image.new('RGBA', image.size)
+	for point in points:
+		point_coords = np.array([[0, 0], point])
+		point_labels = np.array([0, 1])
+		masks, _, _ = predictor.predict(image, point_coords, point_labels)
+		index = 0
+		for mask in masks:
+			maskimage = SegmentAnything.makeMaskImage(mask.T, (0xff, 0xff, 0xff, 0xff))
+			index += 1
+			maskNewImage = SegmentAnything.makeNewImage(image, maskimage)
+			newImage.paste(maskNewImage, (0, 0), maskNewImage)
+			yield newImage
+
+def tab_select(evt: gr.SelectData, state):
+	if evt.target.label == 'point2':
+		state['active'] = 1
+	elif evt.target.label == 'point3':
+		state['active'] = 2
+	else:
+		state['active'] = 0
+	return state
+
+def image_select(evt: gr.SelectData, state, enable1, left1, top1, enable2, left2, top2, enable3, left3, top3):
+	if state['active'] == 2:
+		return [enable1, left1, top1, enable2, left2, top2, True, evt.index[0], evt.index[1]]
+	elif state['active'] == 1:
+		return [enable1, left1, top1, True, evt.index[0], evt.index[1], enable3, left3, top3]
+	return [True, evt.index[0], evt.index[1], enable2, left2, top2, enable3, left3, top3]
+
+with gr.Blocks(title='clip-image') as app:
+	state = gr.State({ 'active': 0 })
+
+	gr.Markdown('''
+	# Clip Image
+	clip an image from given points
+	''')
+	with gr.Row():
+		with gr.Column():
+			image = gr.Image(type='pil')
+			gr.Markdown('click on the image to position')
+			with gr.Tab("point1") as tab1:
+				enable1 = gr.Checkbox(label='enable', value=True)
+				left1 = gr.Slider(maximum=4000, step=1, label='left')
+				top1 = gr.Slider(maximum=4000, step=1, label='top')
+			with gr.Tab("point2") as tab2:
+				enable2 = gr.Checkbox(label='enable')
+				left2 = gr.Slider(maximum=4000, step=1, label='left')
+				top2 = gr.Slider(maximum=4000, step=1, label='top')
+			with gr.Tab("point3") as tab3:
+				enable3 = gr.Checkbox(label='enable')
+				left3 = gr.Slider(maximum=4000, step=1, label='left')
+				top3 = gr.Slider(maximum=4000, step=1, label='top')
+			btnloc = gr.Button(value='check location')
+			with gr.Row():
+				with gr.Column(min_width=160):
+					clearBtn = gr.ClearButton()
+				with gr.Column(min_width=160):
+					btn = gr.Button(value='Submit')
+			inputs = [image, enable1, left1, top1, enable2, left2, top2, enable3, left3, top3]
+		with gr.Column():
+			outputs = [gr.Image(label='segmentation', type='pil')]
+	tab1.select(tab_select, inputs=state, outputs=state)
+	tab2.select(tab_select, inputs=state, outputs=state)
+	tab3.select(tab_select, inputs=state, outputs=state)
+	image.select(image_select, inputs=[state, enable1, left1, top1, enable2, left2, top2, enable3, left3, top3], outputs=[enable1, left1, top1, enable2, left2, top2, enable3, left3, top3])
+	btnloc.click(check_location, inputs=inputs, outputs=outputs)
+	clearBtn.add(inputs + outputs)
+	btn.click(process_image, inputs=inputs, outputs=outputs)
+
+	gr.Examples(
+		[['examples/example1.jpg', True, 200, 250, True, 340, 250, False, 0, 0], ['examples/example2.jpg', True, 256, 256, False, 0, 0, False, 0, 0]],
+		inputs,
+		outputs,
+		process_image,
+		#cache_examples=True,
+	)
+
+app.queue(concurrency_count=5)
+app.launch()

checkpoint/sam_vit_h_4b8939.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e
+size 2564550879

meta_segment_anything.py

@@ -0,0 +1,44 @@
+from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
+import torch
+import numpy as np
+from PIL import Image
+
+class SegmentAnything:
+	def __init__(self):
+		sam_checkpoint = 'checkpoint/sam_vit_h_4b8939.pth'
+		model_type = 'vit_h'
+		sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+		if torch.cuda.is_available():
+			sam.to(device='cuda')
+		self.sam = sam
+
+	def predict(self, image, point_coords, point_labels, box=None):
+		predictor = SamPredictor(self.sam)
+		predictor.set_image(np.array(image, dtype=np.uint8))
+		return predictor.predict(point_coords=point_coords, point_labels=point_labels, box=box)
+
+	def generate(self, image):
+		mask_generator = SamAutomaticMaskGenerator(self.sam)
+		return mask_generator.generate(np.array(image, dtype=np.uint8))
+
+	@staticmethod
+	def makeMaskImage(mask, color):
+		image = Image.new('RGBA', mask.shape)
+		width, height = image.size
+		for x in range(width):
+			for y in range(height):
+				if mask[x, y]:
+					image.putpixel((x, y), color)
+		return image
+
+	@staticmethod
+	def makeNewImage(image, maskImage):
+		newImage = Image.new('RGBA', image.size)
+		timage = maskImage.copy()
+		width, height = timage.size
+		for x in range(width):
+			for y in range(height):
+				_, _, _, a = timage.getpixel((x, y))
+				timage.putpixel((x, y), (0, 0, 0, 255) if a > 0 else (0, 0, 0, 0))
+		newImage.paste(image, (0, 0), timage)
+		return newImage

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+torch
+torchvision
+torchaudio
+git+https://github.com/facebookresearch/segment-anything.git