segment-anything-demo

app.py

@@ -0,0 +1,41 @@
+import gradio as gr
+from meta_segment_anything import SegmentAnything
+from color import Color
+from color_wheel import ColorWheel
+
+def process_image(image):
+	if image is None:
+		yield None
+		return
+
+	color = Color.fromRgb(0xff, 0x0, 0x0)
+	cwt = ColorWheel(color)
+	colors = []
+	for t in cwt.tone15:
+		cw = ColorWheel(Color.fromRgb(t.r, t.g, t.b))
+		for h in cw.hue:
+			colors.append((h.r, h.g, h.b, 0x80))
+
+	generator = SegmentAnything()
+	masks = generator.generate(image)
+	index = 0
+	for mask in masks:
+		if index >= len(colors):
+			break
+		maskimage = SegmentAnything.makeMaskImage(mask['segmentation'].T, colors[index])
+		image.paste(maskimage, (0, 0), maskimage)
+		index += 1
+		yield image
+
+app = gr.Interface(
+	title='Segment Anything Demo',
+	description='generate masks for an entire image',
+	fn=process_image,
+	inputs=gr.Image(type='pil'),
+	outputs=gr.Image(label='segmentation', type='pil'),
+	allow_flagging='never',
+	examples=[['examples/example1.jpg'], ['examples/example2.jpg']],
+	#cache_examples=False
+)
+app.queue(concurrency_count=5)
+app.launch()

color.py

@@ -0,0 +1,36 @@
+class Color:
+	@property
+	def a(self):
+		return self._a
+
+	@property
+	def r(self):
+		return self._r
+
+	@property
+	def g(self):
+		return self._g
+
+	@property
+	def b(self):
+		return self._b
+
+	def __init__(self, a, r, g, b):
+		self._a = a
+		self._r = r
+		self._g = g
+		self._b = b
+
+	@staticmethod
+	def fromArgb(a, r, g, b):
+		return Color(a, r, g, b)
+
+	@staticmethod
+	def fromRgb(r, g, b):
+		return Color(0xff, r, g, b)
+
+	def __eq__(self, other):
+		return self.a == other.a and self.r == other.r and self.g == other.g and self.b == other.b
+
+	def __str__(self):
+		return 'ARGB({},{},{},{})'.format(self.a, self.r, self.g, self.b)

color_wheel.py

@@ -0,0 +1,146 @@
+import math
+from color import Color
+
+class ColorWheel:
+	@property
+	def baseColor(self):
+		return self._baseColor
+
+	@property
+	def hue(self):
+		return [
+			self.baseColor, self.addH(1.0 / 12), self.addH(2.0 / 12), self.addH(3.0 / 12), self.addH(4.0 / 12), self.addH(5.0 / 12),
+			self.addH(-6.0 / 12), self.addH(-5.0 / 12), self.addH(-4.0 / 12), self.addH(-3.0 / 12), self.addH(-2.0 / 12), self.addH(-1.0 / 12)
+		]
+
+	@property
+	def tone(self):
+		return [self.addWhite(-2.0 / 16), self.addWhite(-1.0 / 16), self.baseColor, self.addWhite(1.0 / 16), self.addWhite(2.0 / 16)]
+
+	@property
+	def tone15(self):
+		return [self.addWhite(-7.0 / 16), self.addWhite(-6.0 / 16), self.addWhite(-5.0 / 16), self.addWhite(-4.0 / 16), self.addWhite(-3.0 / 16), 
+				self.addWhite(-2.0 / 16), self.addWhite(-1.0 / 16), self.baseColor, self.addWhite(1.0 / 16), self.addWhite(2.0 / 16),
+				self.addWhite(3.0 / 16), self.addWhite(4.0 / 16), self.addWhite(5.0 / 16), self.addWhite(6.0 / 16), self.addWhite(7.0 / 16)]
+
+	@property
+	def complementaryColors(self):
+		return [self.baseColor, self.addH(0.5)]
+
+	@property
+	def triadicColors(self):
+		return [self.addH(-4.0 / 12), self.baseColor, self.addH(4.0 / 12)]
+
+	@property
+	def splitComplementaryColors(self):
+		return [self.addH(-5.0 / 12), self.baseColor, self.addH(5.0 / 12)]
+
+	@property
+	def analogousColors(self):
+		return [self.addH(-2.0 / 12), self.addH(-1.0 / 12), self.baseColor, self.addH(1.0 / 12), self.addH(2.0 / 12)]
+
+	def __init__(self, c):
+		self._baseColor = c
+		self._r = c.r / 255.0
+		self._g = c.g / 255.0
+		self._b = c.b / 255.0
+
+	@staticmethod
+	def fromHsv(h, s, v):
+		r, g, b = ColorWheel.hsvToRgb(h, s, v)
+		c = Color.fromRgb(round(r * 255), round(g * 255), round(b * 255))
+		return ColorWheel(c)
+
+	def addWhite(self, value):
+		r, g, b = (self._r + value, self._g + value, self._b + value)
+		r = min(max(r, 0.0), 1.0)
+		g = min(max(g, 0.0), 1.0)
+		b = min(max(b, 0.0), 1.0)
+		return self._fromRgb(r, g, b)
+
+	def addH(self, value):
+		h, s, v = ColorWheel.rgbToHsv(self._r, self._g, self._b)
+		h = (h + value) % 1.0
+		if h < 0.0:
+			h += 1.0
+		r, g, b = ColorWheel.hsvToRgb(h, s, v)
+		return self._fromRgb(r, g, b)
+
+	def addS(self, value):
+		h, s, v = ColorWheel.rgbToHsv(self._r, self._g, self._b)
+		s += value
+		s = min(max(s, 0.0), 1.0)
+		r, g, b = ColorWheel.hsvToRgb(h, s, v)
+		return self._fromRgb(r, g, b)
+
+	def addV(self, value):
+		h, s, v = ColorWheel.rgbToHsv(self._r, self._g, self._b)
+		v += value
+		v = min(max(v, 0.0), 1.0)
+		r, g, b = ColorWheel.hsvToRgb(h, s, v)
+		return self._fromRgb(r, g, b)
+
+	@staticmethod
+	def rgbToHsv(r, g, b):
+		if r < 0.0 or r > 1.0:
+			raise ValueError()
+		if g < 0.0 or g > 1.0:
+			raise ValueError()
+		if b < 0.0 or b > 1.0:
+			raise ValueError()
+		cmax = max(r, g, b)
+		cmin = min(r, g, b)
+		h = cmax - cmin
+		if h > 0.0:
+			if cmax == r:
+				h = (g - b) / h
+				if h < 0.0:
+					h += 6.0
+			elif cmax == g:
+				h = 2.0 + (b - r) / h
+			else:
+				h = 4.0 + (r - g) / h
+		h /= 6.0
+		s = cmax - cmin
+		if cmax != 0.0:
+			s /= cmax
+		v = cmax
+		return h, s, v
+
+	@staticmethod
+	def hsvToRgb(h, s, v):
+		if h < 0.0 or h > 1.0:
+			raise ValueError()
+		if s < 0.0 or s > 1.0:
+			raise ValueError()
+		if v < 0.0 or v > 1.0:
+			raise ValueError()
+		r = v
+		g = v
+		b = v
+		if s > 0.0:
+			h *= 6.0
+			i = math.floor(h)
+			f = h - i
+			if i == 1:
+				r *= 1 - s * f
+				b *= 1 - s
+			elif i == 2:
+				r *= 1 - s
+				b *= 1 - s * (1 - f)
+			elif i == 3:
+				r *= 1 - s
+				g *= 1 - s * f
+			elif i == 4:
+				r *= 1 - s * (1 - f)
+				g *= 1 - s
+			elif i == 5:
+				g *= 1 - s
+				b *= 1 - s * f
+			else:
+				g *= 1 - s * (1 - f)
+				b *= 1 - s
+		return r, g, b
+
+	def _fromRgb(self, r, g, b):
+		return Color.fromArgb(self.baseColor.a, round(r * 255), round(g * 255), round(b * 255))

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