print processing time

app.py

@@ -5,6 +5,7 @@ from torch import nn
 import gradio as gr
 import os
 import torch
+import time
 
 feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b5-finetuned-cityscapes-1024-1024")
 model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b5-finetuned-cityscapes-1024-1024")
@@ -50,9 +51,11 @@ def annotation(image:ImageDraw, color_seg:np.array):
             draw.rectangle([(x, y), (x + step_size, y + step_size)], outline="white", width=3)
 
 def call(image): #nparray
+    start = time.time()
 
     print(f"Is CUDA available: {torch.cuda.is_available()}")
-    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+    if (torch.cuda.is_available()):
+        print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
     resized = Image.fromarray(image).resize((1024,1024))
     resized_image = np.array(resized)
@@ -95,6 +98,9 @@ def call(image): #nparray
     out_im_file = Image.fromarray(img)
     annotation(out_im_file, color_seg)
 
+    end = time.time()
+    print(f"processing time: {(end - start):.2f} s")
+
     return out_im_file
 
 # original_image = Image.open("./examples/1.jpg")