Rotation radio buttons

gradio_demo.py

@@ -70,19 +70,7 @@ if torch.cuda.device_count() > 0:
 def check_upload(input_image):
     if input_image is None:
         raise gr.Error("Please provide an image to restore.")
-    return [gr.update(visible = False)] * 2
-
-def rotate_anti_90(image_array):
-    if image_array is None:
-        raise gr.Error("Please provide an image to rotate.")
-
-    return np.array(list(zip(*image_array))[::-1])
-
-def rotate_90(image_array):
-    if image_array is None:
-        raise gr.Error("Please provide an image to rotate.")
-
-    return np.array(list(zip(*image_array[::-1])))
+    return gr.update(visible = True)
 
 def update_seed(is_randomize_seed, seed):
     if is_randomize_seed:
@@ -136,6 +124,7 @@ def llave_process(input_image, temperature, top_p, qs=None):
 
 def stage2_process(
     noisy_image,
+    rotation,
     denoise_image,
     prompt,
     a_prompt,
@@ -165,6 +154,7 @@ def stage2_process(
     allocation
 ):
     print("noisy_image: " + str(noisy_image))
+    print("rotation: " + str(rotation))
     print("denoise_image: " + str(denoise_image))
     print("prompt: " + str(prompt))
     print("a_prompt: " + str(a_prompt))
@@ -213,11 +203,22 @@ def stage2_process(
     else:
         a_prompt = prompt + a_prompt
     print("Final prompt: " + str(a_prompt))
+
     noisy_image = noisy_image if denoise_image is None else denoise_image
+
+    if rotation == 90:
+        noisy_image = np.array(list(zip(*noisy_image[::-1])))
+    elif rotation == 180:
+        noisy_image = np.array(list(zip(*noisy_image[::-1])))
+        noisy_image = np.array(list(zip(*noisy_image[::-1])))
+    elif rotation == -90:
+        noisy_image = np.array(list(zip(*noisy_image))[::-1])
+
     if 1 < downscale:
         input_height, input_width, input_channel = noisy_image.shape
         noisy_image = np.array(Image.fromarray(noisy_image).resize((input_width // downscale, input_height // downscale), Image.LANCZOS))
 
+    # Allocation
     if allocation == 1:
         return restore_in_1min(
             noisy_image, denoise_image, prompt, a_prompt, n_prompt, num_samples, min_size, downscale, upscale, edm_steps, s_stage1, s_stage2, s_cfg, randomize_seed, seed, s_churn, s_noise, color_fix_type, diff_dtype, ae_dtype, gamma_correction, linear_CFG, linear_s_stage2, spt_linear_CFG, spt_linear_s_stage2, model_select, output_format, allocation
@@ -513,11 +514,7 @@ with gr.Blocks(title="SUPIR") as interface:
     gr.HTML(title_html)
 
     input_image = gr.Image(label="Input (*.png, *.webp, *.jpeg, *.gif, *.bmp)", show_label=True, type="numpy", height=600, elem_id="image-input")
-    with gr.Row():
-        with gr.Column():
-            rotate_anti_90_button = gr.Button(value="⤴ Rotate -90°", elem_id="rotate_anti_90_button", visible=False)
-        with gr.Column():
-            rotate_90_button = gr.Button(value="⤵ Rotate +90°", elem_id="rotate_90_button", visible=False)
+    rotation = gr.Radio([["No rotation", 0], ["⤵ Rotate +90°", 90], ["↩ Return 180°", 180], ["⤴ Rotate -90°", -90]], label="Orientation correction", info="Will apply the following rotation before restoring the image; the AI needs a good orientation to understand the content", value=0, interactive=True, visible=False)
     with gr.Group():
         prompt = gr.Textbox(label="Image description", info="Help the AI understand what the image represents; describe as much as possible, especially the details we can't see on the original image; I advise you to write in English because other languages may not be handled", value="", placeholder="A 33 years old man, walking, in the street, Santiago, morning, Summer, photorealistic", lines=3)
         prompt_hint = gr.HTML("You can use a <a href='"'https://huggingface.co/spaces/MaziyarPanahi/llava-llama-3-8b'"'>LlaVa space</a> to auto-generate the description of your image.")
@@ -610,6 +607,7 @@ with gr.Blocks(title="SUPIR") as interface:
         fn = stage2_process,
 	    inputs = [
             input_image,
+            rotation,
             denoise_image,
             prompt,
             a_prompt,
@@ -647,6 +645,7 @@ with gr.Blocks(title="SUPIR") as interface:
         examples = [
                 [
                     "./Examples/Example1.png",
+                    0,
                     None,
                     "Group of people, walking, happy, in the street, photorealistic, 8k, extremely detailled",
                     "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
@@ -677,6 +676,7 @@ with gr.Blocks(title="SUPIR") as interface:
                 ],
                 [
                     "./Examples/Example2.jpeg",
+                    0,
                     None,
                     "The head of a tabby cat, in a house, photorealistic, 8k, extremely detailled",
                     "Cinematic, High Contrast, highly detailed, taken using a Canon EOS R camera, hyper detailed photo - realistic maximum detail, 32k, Color Grading, ultra HD, extreme meticulous detailing, skin pore detailing, hyper sharpness, perfect without deformations.",
@@ -715,20 +715,7 @@ with gr.Blocks(title="SUPIR") as interface:
     input_image.upload(fn = check_upload, inputs = [
         input_image
     ], outputs = [
-        rotate_anti_90_button,
-        rotate_90_button
-    ], queue = False, show_progress = False)
-        
-    rotate_anti_90_button.click(fn = rotate_anti_90, inputs = [
-        input_image
-    ], outputs = [
-        input_image
-    ], queue = False, show_progress = False)
-    
-    rotate_90_button.click(fn = rotate_90, inputs = [
-        input_image
-    ], outputs = [
-        input_image
+        rotation
     ], queue = False, show_progress = False)
 
     denoise_button.click(fn = check, inputs = [
@@ -761,6 +748,7 @@ with gr.Blocks(title="SUPIR") as interface:
         input_image
     ], outputs = [], queue = False, show_progress = False).success(fn=stage2_process, inputs = [
         input_image,
+        rotation,
         denoise_image,
         prompt,
         a_prompt,