from huggingface_hub import hf_hub_download import gradio as gr from fastai.vision.all import * import torch import torchvision.transforms as transforms import numpy as np from PIL import Image repo_id = "hafsa101010/Practica3" model_path = hf_hub_download(repo_id=repo_id, filename="unet.pth") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = torch.jit.load(model_path, map_location=device) model.eval() def transform_image(image): image = image.resize((640, 480)) # Asegurar tamaƱo correcto my_transforms = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) return my_transforms(image).unsqueeze(0).to(device) def predict(img): tensor = transform_image(img) with torch.no_grad(): outputs = model(tensor) outputs = torch.argmax(outputs, 1).cpu().numpy().squeeze() mask = np.zeros_like(outputs, dtype=np.uint8) mask[outputs == 1] = 255 # grape mask[outputs == 2] = 150 # leaves mask[outputs == 3] = 76 # pole mask[outputs == 4] = 74 # pole mask[outputs == 5] = 29 # wood mask[outputs == 6] = 25 # wood return Image.fromarray(mask) gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs=gr.Image(type="pil"), examples=["color_181.jpg", "color_155.jpg"]).launch()