Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| from collections.abc import Iterable | |
| from functools import lru_cache | |
| from pathlib import Path | |
| from PIL import Image, ImageDraw, ImageFont | |
| from .protocol import Primitive | |
| PURPLE = (120, 78, 188) | |
| GREEN = (145, 195, 49) | |
| PINK = (214, 120, 172) | |
| CYAN = (55, 171, 183) | |
| PALETTE = (PURPLE, GREEN, PINK, CYAN) | |
| def _load_font(size: int, multilingual: bool) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: | |
| assets = Path(__file__).resolve().parents[1] / "assets" | |
| if multilingual: | |
| # Pillow does not perform the browser's font fallback. Bundle CJK glyphs | |
| # so Chinese OCR labels render identically on macOS and the Linux Space. | |
| return ImageFont.truetype(str(assets / "NotoSansCJKsc-Regular.otf"), size=size) | |
| candidates = ( | |
| str(assets / "Manrope-variable.ttf"), | |
| "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", | |
| "/usr/share/fonts/truetype/liberation2/LiberationSans-Regular.ttf", | |
| ) | |
| for candidate in candidates: | |
| if Path(candidate).is_file(): | |
| font = ImageFont.truetype(candidate, size=size) | |
| if Path(candidate).name == "Manrope-variable.ttf": | |
| font.set_variation_by_axes([600]) | |
| return font | |
| return ImageFont.load_default() | |
| def _font(size: int, text: str = "") -> ImageFont.FreeTypeFont | ImageFont.ImageFont: | |
| return _load_font(size, not text.isascii()) | |
| def _xy(values: tuple[int, ...], width: int, height: int) -> tuple[int, ...]: | |
| if len(values) == 2: | |
| return round(values[0] * (width - 1) / 999), round(values[1] * (height - 1) / 999) | |
| return ( | |
| round(values[0] * (width - 1) / 999), | |
| round(values[1] * (height - 1) / 999), | |
| round(values[2] * (width - 1) / 999), | |
| round(values[3] * (height - 1) / 999), | |
| ) | |
| def _intersects(a, b, gap=2): | |
| return a[0] < b[2]+gap and b[0] < a[2]+gap and a[1] < b[3]+gap and b[1] < a[3]+gap | |
| def layout_box_labels(primitives, width, height, font, padding=4): | |
| """One collision-free badge per exact category; returned boxes stay untouched.""" | |
| groups = {} | |
| for primitive in primitives: | |
| if len(primitive.coordinates) == 4: | |
| groups.setdefault(primitive.label, []).append(_xy(primitive.coordinates, width, height)) | |
| area = lambda b: max(0,b[2]-b[0])*max(0,b[3]-b[1]) | |
| small = [b for boxes in groups.values() for b in boxes if area(b)<width*height*.01] | |
| measure = ImageDraw.Draw(Image.new("RGB", (1, 1))) | |
| occupied, placed = [], [] | |
| # Give small targets first choice of nearby space. Large regions have more | |
| # candidate anchors. Ordering affects label placement only, never predictions. | |
| for label, boxes in sorted(groups.items(), key=lambda item: max(map(area,item[1]))): | |
| label_font = _font(getattr(font, "size", 13), label) if not label.isascii() else font | |
| text = label if len(label)<=38 else label[:37]+"…" | |
| while text: | |
| bounds = measure.textbbox((0,0),text,font=label_font) | |
| tw, th = bounds[2]-bounds[0]+2*padding, bounds[3]-bounds[1]+2*padding | |
| if tw<=width-2 and th<=height-2: break | |
| text = text[:-2]+"…" if len(text)>2 else "" | |
| if not text: continue | |
| candidates=[] | |
| for target in sorted(boxes,key=area,reverse=True): | |
| x1,y1,x2,y2=target | |
| for x,y in [(x1,y1-th-2),(x1,y1+2),(x1,y2+2),(x2-tw,y1-th-2), | |
| (x2-tw,y2-th-2),(x2+2,y1),(x1-tw-2,y1)]: | |
| x=max(1,min(x,width-tw-1));y=max(1,min(y,height-th-1)) | |
| candidates.append(((x,y,x+tw,y+th),target)) | |
| chosen=next(((b,t) for b,t in candidates if not any(_intersects(b,o) for o in occupied+small)),None) | |
| if chosen is None: | |
| # Rare crowded anchors use the nearest free area with a leader line. | |
| # Never put a badge on another badge or cover a small prediction. | |
| free=[] | |
| for y in range(1,max(2,height-th),max(4,th//2)): | |
| for x in range(1,max(2,width-tw),max(6,tw//4)): | |
| b=(x,y,x+tw,y+th) | |
| if any(_intersects(b,o) for o in occupied+small): continue | |
| for t in boxes: | |
| dx=max(t[0]-b[2],b[0]-t[2],0);dy=max(t[1]-b[3],b[1]-t[3],0) | |
| # Prefer top-left anchors among equally near locations. | |
| score=dx*dx+dy*dy+.01*((x-t[0])**2+(y-t[1])**2) | |
| free.append((score,b,t)) | |
| if free: | |
| _,b,t=min(free,key=lambda item:item[0]);chosen=(b,t) | |
| if chosen: | |
| b,target=chosen;occupied.append(b) | |
| placed.append(dict(label=label,text=text,box=b,target=target,text_bounds=bounds,font=label_font)) | |
| return placed | |
| def draw_primitives( | |
| image: Image.Image, | |
| primitives: Iterable[Primitive], | |
| references: Iterable[Primitive] = (), | |
| title: str | None = None, | |
| ) -> Image.Image: | |
| canvas = image.convert("RGBA") | |
| overlay = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) | |
| draw = ImageDraw.Draw(overlay) | |
| width, height = canvas.size | |
| primitives = list(primitives) | |
| dense = len(primitives) > 12 | |
| # Fine boundaries keep a dense scene readable; short corners make small | |
| # boxes discoverable without filling, moving, or removing any prediction. | |
| line = max(1 if dense else 2, round(min(width, height) / (420 if dense else 260))) | |
| radius = max(6, line * 3) | |
| font = _font(max(12 if dense else 13, round(min(width, height) / 48))) | |
| label_colors: dict[str, tuple[int, int, int]] = {} | |
| for ref in references: | |
| if len(ref.coordinates) != 4: | |
| continue | |
| box = _xy(ref.coordinates, width, height) | |
| draw.rectangle(box, outline=(*GREEN, 255), width=line * 2) | |
| x1, y1, x2, y2 = box | |
| dash = max(5, line * 3) | |
| for x in range(x1, x2, dash * 2): | |
| draw.line((x, y1, min(x + dash, x2), y1), fill=(*GREEN, 255), width=line) | |
| draw.line((x, y2, min(x + dash, x2), y2), fill=(*GREEN, 255), width=line) | |
| # Opaque RGBA writes replace earlier pixels in a Pillow layer. In particular, | |
| # a later lake/cloud fill used to erase every small object drawn inside it. | |
| # Keep region interiors untouched; geometry and labels are separate passes. | |
| for primitive in primitives: | |
| color = label_colors.setdefault(primitive.label, PALETTE[len(label_colors) % len(PALETTE)]) | |
| coords = _xy(primitive.coordinates, width, height) | |
| if len(coords) == 2: | |
| x, y = coords | |
| draw.ellipse( | |
| (x - radius, y - radius, x + radius, y + radius), | |
| outline=(*color, 255), | |
| width=line, | |
| ) | |
| draw.line((x - radius * 2, y, x + radius * 2, y), fill=(*color, 255), width=line) | |
| draw.line((x, y - radius * 2, x, y + radius * 2), fill=(*color, 255), width=line) | |
| else: | |
| x1, y1, x2, y2 = coords | |
| draw.rectangle(coords, outline=(*color, 255), width=line) | |
| if dense: | |
| corner = min(max(4, round(min(width, height) / 80)), (x2-x1)//3, (y2-y1)//3) | |
| if corner > 0: | |
| for x, y, dx, dy in ((x1,y1,1,1),(x2,y1,-1,1),(x1,y2,1,-1),(x2,y2,-1,-1)): | |
| draw.line((x+dx*corner,y,x,y,x,y+dy*corner), fill=(*color,255), width=max(2,line)) | |
| # Labels remain visible in dense scenes, once per category across all boxes. | |
| padding=max(4,line*2) | |
| for badge in layout_box_labels(primitives,width,height,font,padding): | |
| text,box,target=badge["text"],badge["box"],badge["target"] | |
| color=label_colors[badge["label"]] | |
| tx,ty=box[:2];tb=badge["text_bounds"] | |
| # Attach displaced badges to their actual prediction with a fine leader. | |
| if not _intersects(box,target,gap=4): | |
| cx,cy=(box[0]+box[2])/2,(box[1]+box[3])/2 | |
| ex=max(target[0],min(cx,target[2]));ey=max(target[1],min(cy,target[3])) | |
| sx=max(box[0],min(ex,box[2]));sy=max(box[1],min(ey,box[3])) | |
| draw.line((sx,sy,ex,ey),fill=(*color,230),width=1) | |
| light=tuple(round(c*.10+255*.90) for c in color) | |
| draw.rounded_rectangle(box,radius=3,fill=(*light,248),outline=(*color,255),width=1) | |
| draw.text((tx+padding-tb[0],ty+padding-tb[1]),text,font=badge["font"],fill=(38,30,49,255)) | |
| if title: | |
| title_font = _font(getattr(font, "size", 13), title) | |
| tb = draw.textbbox((0, 0), title, font=title_font) | |
| draw.rounded_rectangle( | |
| (10, 10, min(width - 10, tb[2] + 28), tb[3] + 24), | |
| radius=10, | |
| fill=(255, 255, 255, 225), | |
| outline=(*PURPLE, 255), | |
| width=line, | |
| ) | |
| draw.text((18, 16), title, font=title_font, fill=(*PURPLE, 255)) | |
| return Image.alpha_composite(canvas, overlay).convert("RGB") | |