Spaces:
Sleeping
Sleeping
import os | |
import gradio as gr | |
import pytesseract | |
import yolov5 | |
from transformers import CLIPProcessor, CLIPModel | |
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") | |
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") | |
# # load model | |
# model = yolov5.load('keremberke/yolov5m-license-plate') | |
# # set model parameters | |
# model.conf = 0.5 # NMS confidence threshold | |
# model.iou = 0.25 # NMS IoU threshold | |
# model.agnostic = False # NMS class-agnostic | |
# model.multi_label = False # NMS multiple labels per box | |
# model.max_det = 1000 # maximum number of detections per image | |
def license_plate_detect(img): | |
results = model(img, size=640) | |
# parse results | |
predictions = results.pred[0] | |
if len(predictions): | |
boxes = predictions[:, :4] # x1, y1, x2, y2 | |
return boxes | |
def read_license_number(img): | |
boxes = license_plate_detect(img) | |
if boxes is not None: | |
return [pytesseract.image_to_string( | |
img.crop(bbox.tolist())) | |
for bbox in boxes] | |
def zero_shot_classification(image, labels): | |
print(type(image)) | |
inputs = processor(text=labels, | |
images=image, | |
return_tensors="pt", | |
padding=True) | |
print(type(inputs)) | |
print(inputs) | |
outputs = model(**inputs) | |
logits_per_image = outputs.logits_per_image # this is the image-text similarity score | |
return logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities | |
installed_list = [] | |
# image = Image.open(requests.get(url, stream=True).raw) | |
def check_solarplant_installed_by_license(license_number_list): | |
if len(installed_list): | |
return [license_number in installed_list | |
for license_number in license_number_list] | |
def check_solarplant_installed_by_image(image, output_label=False): | |
zero_shot_class_labels = ["bus with solar panel grids", | |
"bus without solar panel grids"] | |
probs = zero_shot_classification(image, zero_shot_class_labels) | |
if output_label: | |
return zero_shot_class_labels[probs.argmax().item()] | |
return probs.argmax().item() == 0 | |
def check_solarplant_broken(image): | |
zero_shot_class_labels = ["white broken solar panel", | |
"normal black solar panel grids"] | |
probs = zero_shot_classification(image, zero_shot_class_labels) | |
idx = probs.argmax().item() | |
return zero_shot_class_labels[idx][1-idx] | |
def greet(img): | |
print(type(img)) | |
# lns = read_license_number(img) | |
lns = ['1','2','3'] | |
if len(lns): | |
# return (seg, | |
return ("θ»ηοΌ " + '; '.join(lns) + "\n\n" \ | |
+ "ι‘εοΌ "+ check_solarplant_installed_by_image(img, True) + "\n\n" \ | |
+ "ηζ οΌ" + check_solarplant_broken(img)) | |
return (img, "η©Ίε°γγγ") | |
iface = gr.Interface(fn=greet, inputs=gr.Image(type="pil"), outputs="text") | |
iface.launch() |