Update app.py

app.py

@@ -1,3 +1,8 @@
+"""
+Created By: ishwor subedi
+Date: 2024-05-19
+"""
+
 from PIL import Image
 import torch
 import torch.nn.functional as F
@@ -62,8 +67,16 @@ def example_inference(image):
 
 
 original_image, binary_image = None, None
-paths = [os.path.join("bg_images", file) for file in os.listdir("bg_images")]
-images = [Image.open(path) for path in paths]
+colors = [Image.open(path) for path in
+          [os.path.join("bg_images/color", file) for file in os.listdir("bg_images/color")]]
+houses = [Image.open(path) for path in
+          [os.path.join("bg_images/house", file) for file in os.listdir("bg_images/house")]]
+natures = [Image.open(path) for path in
+           [os.path.join("bg_images/nature", file) for file in os.listdir("bg_images/nature")]]
+studios = [Image.open(path) for path in
+           [os.path.join("bg_images/studio", file) for file in os.listdir("bg_images/studio")]]
+walls = [Image.open(path) for path in [os.path.join("bg_images/wall", file) for file in os.listdir("bg_images/wall")]]
+woods = [Image.open(path) for path in [os.path.join("bg_images/wood", file) for file in os.listdir("bg_images/wood")]]
 
 with gr.Blocks(
         theme=gr.themes.Default(primary_hue=gr.themes.colors.red, secondary_hue=gr.themes.colors.indigo)) as demo:
@@ -87,7 +100,18 @@ with gr.Blocks(
 
 
     with gr.Row():
-        examples = gr.Examples(examples=images, inputs=[hidden_img])
+        examples = gr.Examples(examples=studios, inputs=[hidden_img], label="Studio Backgrounds")
+
+    with gr.Row():
+        examples6 = gr.Examples(examples=colors, inputs=[hidden_img], label="Color Backgrounds")
+
+    with gr.Row():
+        examples2 = gr.Examples(examples=walls, inputs=[hidden_img], label="Wall Backgrounds")
+        examples3 = gr.Examples(examples=natures, inputs=[hidden_img], label="Nature Backgrounds")
+
+    with gr.Row():
+        examples4 = gr.Examples(examples=houses, inputs=[hidden_img], label="House Backgrounds")
+        examples5 = gr.Examples(examples=woods, inputs=[hidden_img], label="Wood Backgrounds")
 
     with gr.Row():
         submit = gr.Button("Submit")
@@ -96,7 +120,7 @@ with gr.Blocks(
 
     def generate_img(image, background):
         orig_img = example_inference(image)
-        height, width = orig_img.size
+        width, height = orig_img.size
 
         background = Image.fromarray(background).resize((width, height))
         orig_img = Image.fromarray(np.array(orig_img)).resize((width, height))