Spaces:
Sleeping
Sleeping
| from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation | |
| from PIL import Image, ImageDraw | |
| import numpy as np | |
| from torch import nn | |
| import gradio as gr | |
| import os | |
| feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b5-finetuned-cityscapes-1024-1024") | |
| model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b5-finetuned-cityscapes-1024-1024") | |
| def cityscapes_palette(): | |
| """Cityscapes palette for external use.""" | |
| return [[128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156], | |
| [190, 153, 153], [153, 153, 153], [250, 170, 30], [220, 220, 0], | |
| [107, 142, 35], [152, 251, 152], [70, 130, 180], [220, 20, 60], | |
| [255, 0, 0], [0, 0, 142], [0, 0, 70], [0, 60, 100], [0, 80, 100], | |
| [0, 0, 230], [119, 11, 32]] | |
| def cityscapes_classes(): | |
| """Cityscapes class names for external use.""" | |
| return [ | |
| 'road', 'sidewalk', 'building', 'wall', 'fence', 'pole', | |
| 'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky', | |
| 'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', | |
| 'bicycle' | |
| ] | |
| def annotation(image:ImageDraw, color_seg:np.array): | |
| assert image.size == (1024, 1024) | |
| assert color_seg.shape == (1024, 1024, 3) | |
| blocks = 4 # 4x4 sub grid | |
| step_size = 256 # sub square edge size | |
| draw = ImageDraw.Draw(image) | |
| sub_square_xy = [(x,y) for x in range(0, blocks * step_size, step_size) for y in range(0, blocks * step_size, step_size)] | |
| # print(f"{sub_square_xy=}") | |
| for (x,y) in sub_square_xy: | |
| reduced_seg = color_seg.sum(axis=2) # collapsing all colors into 1024 x 1024 | |
| # print(f"{reduced_seg.shape=}") | |
| sub_square_seg = reduced_seg[ y:y+step_size, x:x+step_size] | |
| # print(f"{sub_square_seg.shape=}, {sub_square_seg.sum()}") | |
| if (sub_square_seg.sum() > 1000000): | |
| print("light found at square ", x, y) | |
| draw.rectangle([(x, y), (x + step_size, y + step_size)], outline=128, width=3) | |
| def call(image: Image): | |
| resized_image = original_image.resize((1024,1024)) | |
| print(f"{np.array(resized_image).shape=}") # 1024, 1024, 3 | |
| inputs = feature_extractor(images=resized_image, return_tensors="pt") | |
| outputs = model(**inputs) | |
| print(f"{outputs.logits.shape=}") # shape (batch_size, num_labels, height/4, width/4) -> 3, 19, 256 ,256 | |
| # print(f"{logits}") | |
| # First, rescale logits to original image size | |
| interpolated_logits = nn.functional.interpolate( | |
| outputs.logits, | |
| size=resized_image.size[::-1], # (height, width) | |
| mode='bilinear', | |
| align_corners=False) | |
| print(f"{interpolated_logits.shape=}, {outputs.logits.shape=}") # 1, 19, 1024, 1024 | |
| # Second, apply argmax on the class dimension | |
| seg = interpolated_logits.argmax(dim=1)[0] | |
| print(f"{seg.shape=}") | |
| color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3 | |
| print(f"{color_seg.shape=}") | |
| for label, color in enumerate(cityscapes_palette()): | |
| if (label == 6): color_seg[seg == label, :] = color | |
| # Convert to BGR | |
| color_seg = color_seg[..., ::-1] | |
| print(f"{color_seg.shape=}") | |
| # Show image + mask | |
| img = np.array(resized_image) * 0.5 + color_seg * 0.5 | |
| img = img.astype(np.uint8) | |
| out_im_file = Image.fromarray(img) | |
| annotation(out_im_file, color_seg) | |
| return out_im_file | |
| original_image = Image.open("./examples/1.jpg") | |
| print(f"{np.array(original_image).shape=}") # eg 729, 1000, 3 | |
| # out = call(original_image) | |
| # out.save("out2.jpeg") | |
| title = "Traffic Light Detector" | |
| description = "Experiment traffic light detection to evaluate the value of captcha security controls" | |
| iface = gr.Interface(fn=call, | |
| inputs="image", | |
| outputs="image", | |
| title=title, | |
| description=description, | |
| examples=[ | |
| os.path.join(os.path.dirname(__file__), "examples/1.jpg"), | |
| os.path.join(os.path.dirname(__file__), "examples/2.jpg") | |
| ], | |
| thumbnail="thumbnail.webp") | |
| iface.launch() |