Update app.py

app.py

@@ -1,160 +1,29 @@
-DEVELOP_MODE = False 
-USER_MODE = not DEVELOP_MODE
-AZURE_SEARCH_KEY = ""
-
-import os
-from pathlib import Path
 import gradio as gr
-from fastai.vision.all import *
-
-
-if DEVELOP_MODE:
-    import fastbook
-    from fastbook import *
-    from fastai.vision.widgets import *
-    from fastai.vision.all import *
-    fastbook.setup_book()
-
-    import uuid
-    import requests
-    import imghdr
-    from PIL import Image
-    import numpy as np
-    
-
-attn_slicing_enabled = True
-
-
-def download_unique_image(url, folder_path):
-    try:
-        response = requests.get(url, timeout=10)
-        content_type = response.headers.get('Content-Type')
-        if content_type.startswith('image'):
-            image_type = imghdr.what(None, response.content)
-            if image_type == 'jpeg':
-                extension = 'jpg'
-            else:
-                extension = image_type
-            filename = str(uuid.uuid4()) + '.' + extension
-            filepath = os.path.join(folder_path, filename)
-            with open(filepath, 'wb') as f:
-                f.write(response.content)
-    except:
-        pass
-
-
-def remove_corrupted_images(folder_path):
-    count = 0
-    for file_name in os.listdir(folder_path):
-        file_path = os.path.join(folder_path, file_name)
-        try:
-            with Image.open(file_path) as img:
-                pass
-        except Exception as err:
-            os.remove(file_path)
-            count += 1
-
-
-def normalize_dog_name(dog_name):
-    return dog_name.replace(' ', '_').lower()
-
-
-def download_images_():
-    dogs = {
-        'Zwergspitz Dog': [],
-        'Bouledogue Français Dog': [],
-        'Shih Tzu Dog': [],
-        'Rottweiler Dog': [],
-        'Pug Dog': [],
-        'Golden Retriever Dog': [],
-        'Deutscher Schäferhund Dog': [],
-        'Yorkshire Terrier Dog': [],
-        'Border Collie Dog': [],
-        'Dachshund Dog': [],
-        'Poodle Dog': [],
-        'Labrador Retriever Dog': [],
-        'Pinscher Dog': [],
-        'Golden Retriever': [],
+from transformers import pipeline
+
+dogs = {
+    'Zwergspitz Dog': [],
+    'Bouledogue Français Dog': [],
+    'Shih Tzu Dog': [],
+    'Rottweiler Dog': [],
+    'Pug Dog': [],
+    'Golden Retriever Dog': [],
+    'Deutscher Schäferhund Dog': [],
+    'Yorkshire Terrier Dog': [],
+    'Border Collie Dog': [],
+    'Dachshund Dog': [],
+    'Poodle Dog': [],
+    'Labrador Retriever Dog': [],
+    'Pinscher Dog': [],
+    'Golden Retriever': [],
     }
-    DOGS_NAMES = tuple(dogs.keys())
-    if DEVELOP_MODE:
-        if not PATH.exists():
-            PATH.mkdir()
-            for dog_name in DOGS_NAMES:
-                urls = search_images_bing(
-                    AZURE_KEY, dog_name).attrgot('contentUrl')
-                dogs[dog_name] = urls
-
-                dest = os.path.join(PATH, normalize_dog_name(dog_name))
-                if not os.path.exists(dest):
-                    os.mkdir(dest)
-                download_images(dest, urls=urls)
-                remove_corrupted_images(dest)
-    return [dog.replace('Dog', '') for dog in DOGS_NAMES]
 
+pipeline = pipeline(task="image-classification", model="ericxlima/DogsClassifierModel")
 
-def train_model():
-    dogs_datablock = DataBlock(
-        blocks=(ImageBlock, CategoryBlock),
-        get_items=get_image_files,
-        splitter=RandomSplitter(valid_pct=0.2, seed=42),
-        get_y=parent_label,
-        item_tfms=[Resize(128, ResizeMethod.Squish),
-                   Resize(128, ResizeMethod.Pad, pad_mode='zeros'),
-                   RandomResizedCrop(128, min_scale=0.3),
-                   ]
-    )
-    dogs_dataloaders = dogs_datablock.dataloaders(PATH)
-    # dogs_dataloaders = dogs_dataloaders.new(
-    #     item_tfms=Resize(128, ResizeMethod.Squish))
-    learn_ = vision_learner(dogs_dataloaders, resnet18, metrics=error_rate)
-    learn_.fine_tune(4)
-    learn_.export('dogs.pkl')
-    return learn_
-
-
-def classify_image(image):
-    global learing
-    pred, pred_idx, probs = learing.predict(image)
-    return f"Prediction: {pred.replace('_', '').replace('dog', '').title()};\n Probability: {probs[pred_idx]:.04f}"
-
-
-def get_model_():
-    path = Path()
-    model = None
-
-    if any(file.endswith('.pkl') for file in os.listdir(path)):
-        model_ = load_learner('dogs.pkl')
-    else:
-        model_ = train_model()
-    return model_
-
-
-AZURE_KEY = os.environ.get(
-    'AZURE_SEARCH_KEY',
-    AZURE_SEARCH_KEY,
-)
-PATH = Path('dogs')
-
-dogs = download_images_()
-learing = get_model_()
-
-
-# Gradio
-iface = gr.Interface(
-    classify_image,
-    inputs="image",
-    outputs="text",
-    title="Classificação de Imagens",
-    description="Insira uma imagem para ser classificada"
-)
-
+def predict(image):
+  predictions = pipeline(image)
+  return {p["label"]: p["score"] for p in predictions}
 
-def set_mem_optimizations(pipe):
-    if attn_slicing_enabled:
-        pipe.enable_attention_slicing()
-    else:
-        pipe.disable_attention_slicing()
 
 
 def list_breeds():
@@ -162,7 +31,7 @@ def list_breeds():
     html = "<div class='row'>"
     html += "<div class='column'>"
     html += "<h2>List of breed dogs trained:</h2>"
-    html += "<ol>" + "".join([f"<li>{breed}</li>" for breed in dogs]) + "</ol>"
+    html += "<ol>" + "".join([f"<li>{breed}</li>" for breed in list(dogs.keys())]) + "</ol>"
     html += "</div>"
     html += "<div class='column'>"
     html += "<h2>Author:</h2>"
@@ -175,16 +44,16 @@ def list_breeds():
 
 image = gr.Image(shape=(224, 224))
 label = gr.Label(num_top_classes=3)
-breeds_list = list_breeds()
+# breeds_list = list_breeds()
 
 demo = gr.Interface(
-    fn=classify_image,
+    fn=predict,
     inputs=image,
     outputs=label,
     title="🐶 Dog Breed Classifier",
     interpretation="default",
     description="Upload an image of a dog and the model will predict its breed.",
-    article=breeds_list,
+    # article=breeds_list,
     css=".row { display: flex; } .column { flex: 50%; }",
 )