File size: 2,648 Bytes
53d158f e525f34 53d158f e525f34 53d158f e525f34 38a0343 e525f34 53d158f e525f34 53d158f e525f34 53d158f e525f34 53d158f e525f34 53d158f a2cc3e1 53d158f e525f34 53d158f e525f34 53d158f e525f34 53d158f e525f34 53d158f e525f34 53d158f |
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 79 80 81 82 83 |
#!/usr/bin/env python
# coding: utf-8
from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
from huggingface_hub import snapshot_download
import os
from torchvision import transforms
classes = {0: "Defective", 1: "Good"}
#model_path = "best_int8_openvino_model"
model_path = "./best.pt"
def load_model_local():
detection_model = YOLO(model_path, task='classify') # Load the model
return detection_model
def load_model(repo_id):
download_dir = snapshot_download(repo_id)
print(download_dir)
path = os.path.join(download_dir, "best_int8_openvino_model")
print(path)
detection_model = YOLO(path, task='classify')
return detection_model
def predict(pilimg):
source = pilimg
# Call the model to transform image size
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
source = transform(source) # Update the source image size to 224x224, 1 of 2 sizes accepted by Yolo classification model
#result = detection_model.predict(source, conf=0.5, iou=0.6)
result = detection_model.predict(source) # Make prediction
# Get the top prediction
label = result[0].probs.top1
class_names = detection_model.names # Retrieves the class names mapping (dict-like)
classified_type = class_names[label] # Map numeric label to class name
print (">>> Class : ", classified_type)
confidence = result[0].probs.top1conf # Get the top class confidence
print(">>> Confidence : ", confidence)
annotated_image = pilimg.convert("RGB")
draw = ImageDraw.Draw(annotated_image)
font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
font = ImageFont.truetype(font_path, 30)
#font = ImageFont.truetype("font/arialbd.ttf", 30) # Use arial.ttf for bold font
if classified_type == classes[0]:
draw.text((300, 10), classified_type, fill="red", font=font)
gr.Warning("Defect detected, BAD!.")
else:
draw.text((300, 10), classified_type, fill="green", font=font)
gr.Info("No defect detected,GOOD!")
#draw.text((300, 10), classified_type, fill="red", font=font)
return annotated_image
detection_model = load_model_local()
title = "Detect the status of the cap, DEFECTIVE or GOOD"
interface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil", label="Input Image"),
outputs=gr.Image(type="pil", label="Classification result"),
title=title,
)
# Launch the interface
interface.launch(share=True)
|