File size: 2,690 Bytes
76c920d cd4ffa3 f2943ab cd4ffa3 f2943ab 76c920d cd4ffa3 f2943ab cd4ffa3 a8bc6ca a09c55b f2943ab a09c55b 5533283 af4145f 76c920d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import gradio as gr
import torch
from fastai.vision.all import *
from PIL import ImageFilter, ImageEnhance, ImageDraw
from diffusers.utils import make_image_grid
from tqdm import tqdm
from diffusers import AutoPipelineForInpainting, LCMScheduler, DDIMScheduler
from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel
import numpy as np
from PIL import Image
from datetime import datetime
preferred_dtype = torch.float32
preferred_device = "cuda" if torch.cuda.is_available() else "cpu"
def label_func(fn): return path/"labels"/f"{fn.stem}_P{fn.suffix}"
segmodel = load_learner("camvid-512.pkl")
working_size = (512, 512)
seg_vocabulary = ['Animal', 'Archway', 'Bicyclist', 'Bridge', 'Building', 'Car',
'CartLuggagePram', 'Child', 'Column_Pole', 'Fence', 'LaneMkgsDriv',
'LaneMkgsNonDriv', 'Misc_Text', 'MotorcycleScooter', 'OtherMoving',
'ParkingBlock', 'Pedestrian', 'Road', 'RoadShoulder', 'Sidewalk',
'SignSymbol', 'Sky', 'SUVPickupTruck', 'TrafficCone',
'TrafficLight', 'Train', 'Tree', 'Truck_Bus', 'Tunnel',
'VegetationMisc', 'Void', 'Wall']
ban_cars_mask = np.array([0, 0, 0, 0, 0, 1,
0, 0, 1, 0, 1,
1, 1, 0, 0,
1, 0, 1, 1, 1,
1, 0, 1, 1,
1, 0, 0, 0, 1,
0, 1, 0], dtype=np.uint8)
def get_seg_mask(img):
mask = segmodel.predict(img)[0]
return mask
def display_mask(img, mask):
# Convert the grayscale mask to RGB
mask_rgb = np.stack([np.zeros_like(mask), mask, np.zeros_like(mask)], axis=-1)
# Convert the image to PIL format
img_pil = Image.fromarray(img)
# Convert the mask to PIL format
mask_pil = Image.fromarray((mask_rgb * 255).astype(np.uint8))
# Overlay the mask on the image
overlaid_img = Image.blend(img_pil, mask_pil, alpha=0.5)
return overlaid_img
def redact_image(img):
img = img.resize((256, 256))
mask = get_seg_mask(img)
car_mask = ban_cars_mask[mask]
return display_mask(img, car_mask)
def flip(img):
return np.flipud(img)
def app(img):
start_time = datetime.now().timestamp()
img = img.resize(working_size)
mask = ban_cars_mask[get_seg_mask(img)]
overlay_img = Image.fromarray(np.stack([img[:, :, 0], mask / 2, img[:,:,2]], axis=-1))
end_time = datetime.now().timestamp()
draw = ImageDraw.Draw(overlay_img)
draw.text((10, 10), f"Duration: {int(1000 * (end_time - start_time))}ms", fill=(255, 255, 255))
return overlay_img
#ideally:
#iface = gr.Interface(app, gr.Image(sources=["webcam"], streaming=True), "image", live=True)
iface = gr.Interface(app, gr.Image(), "image")
iface.launch()
|