Create app.py

app.py

@@ -0,0 +1,63 @@
+import gradio as gr
+import torch
+from PIL import Image
+from torchvision import transforms
+from transformers import T5Tokenizer, ViTFeatureExtractor
+
+# Model loading and setting up the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = torch.load("model_vit_ai.pt", map_location=device)
+model.to(device)
+
+# Tokenizer and Feature Extractor
+tokenizer = T5Tokenizer.from_pretrained('t5-base')
+feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
+
+# Define the image preprocessing
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=feature_extractor.image_mean, std=feature_extractor.image_std)
+])
+
+def preprocess_image(image):
+    image = Image.fromarray(image.astype('uint8'), 'RGB')
+    image = transform(image)
+    return image.unsqueeze(0)
+
+def generate_caption(image):
+    model.eval()
+    with torch.no_grad():
+        image_tensor = preprocess_image(image).to(device)
+        decoder_input_ids = torch.full((1, 1), model.decoder_start_token_id, dtype=torch.long, device=device)
+
+        for _ in range(50):
+            outputs = model(images=image_tensor, decoder_ids=decoder_input_ids)
+            next_token_logits = outputs.logits[:, -1, :]
+            next_token_id = next_token_logits.argmax(1, keepdim=True)
+            decoder_input_ids = torch.cat([decoder_input_ids, next_token_id], dim=-1)
+
+            if torch.eq(next_token_id, tokenizer.eos_token_id).all():
+                break
+
+        caption = tokenizer.decode(decoder_input_ids.squeeze(0), skip_special_tokens=True)
+    return caption
+
+sample_images = [
+    "sample_image1.jpg",
+    "sample_image2.jpg",
+    "sample_image3.jpg"
+]
+
+# Define Gradio interface
+interface = gr.Interface(
+    fn=generate_caption,
+    inputs=gr.inputs.Image(source="upload", tool='editor', type="numpy", label="Upload an image or take a photo"),
+    outputs='text',
+    examples=sample_images,
+    title="Image Captioning Model",
+    description="Upload an image, select a sample image, or use your webcam to take a photo and generate a caption."
+)
+
+# Run the interface
+interface.launch(debug=True)