Update medical_pipeline.py

medical_pipeline.py

@@ -91,109 +91,10 @@ class MedicalPipeline:
             num_inference_steps=50,
             **kwargs,
     ):
-        if organ is None:
-            if not isinstance(keys, List):
-                if keys is None:
-                    idx = random.randint(1, 7)
-                    if idx == 1:
-                        keys = 'AMOS2022'
-                        organ = 'abdomen CT scans'
-                        kind = self.get_random_values(self.AMOS2022)
-                        img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
-                        mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
-                    elif idx == 2:
-                        keys = 'BUSI'
-                        organ = 'breast ultrasound'
-                        choice = random.random()
-                        if choice < 0.5:
-                            kind = 'normal'
-                        else:
-                            kind = 'breast tumor'
-                        img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
-                        mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
-                    elif idx == 3:
-                        keys = 'ACDC'
-                        organ = 'cardiovascular ventricle mri'
-                        kind = self.get_random_values(self.ACDC)
-                        img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
-                        mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
-                    elif idx == 4:
-                        keys = 'CVC-ClinicDB'
-                        organ = 'polyp colonoscopy'
-                        kind = 'polyp'
-                        img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
-                        mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
-                    elif idx == 5:
-                        keys = 'kvasir-seg'
-                        organ = 'polyp colonoscopy'
-                        kind = 'polyp'
-                        img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
-                        mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
-                    elif idx == 6:
-                        keys = 'LiTS2017'
-                        organ = 'abdomen CT scans'
-                        kind = self.get_random_values(self.LiTS2017)
-                        img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
-                        mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
-                    elif idx == 7:
-                        keys = 'KiTS2019'
-                        organ = 'abdomen CT scans'
-                        kind = self.get_random_values(self.KiTS2019)
-                        img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
-                        mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
-                    else:
-                        raise RuntimeError('no mode')
 
-                    print(img_prompt, mask_prompt)
-
-                else:
-                    if keys == 'AMOS2022':
-                        organ = 'abdomen CT scans'
-                        kind = self.get_random_values(self.AMOS2022)
-                    elif keys == 'BUSI':
-                        organ = 'breast ultrasound'
-                        choice = random.random()
-                        if choice < 0.5:
-                            kind = 'normal'
-                        else:
-                            kind = 'breast tumor'
-                    elif keys == 'ACDC':
-                        organ = 'cardiovascular ventricle mri'
-                        kind = self.get_random_values(self.ACDC)
-                    elif keys == 'CVC-ClinicDB':
-                        organ = 'polyp colonoscopy'
-                        kind = 'polyp'
-                    elif keys == 'kvasir-seg':
-                        organ = 'polyp colonoscopy'
-                        kind = 'polyp'
-                    elif keys == 'LiTS2017':
-                        organ = 'abdomen CT scans'
-                        kind = self.get_random_values(self.LiTS2017)
-                    elif keys == 'KiTS2019':
-                        organ = 'abdomen CT scans'
-                        kind = self.get_random_values(self.KiTS2019)
-                    else:
-                        raise RuntimeError('undefined keys')
-
-                    img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
-                    mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
-                    keys = keys
-                    print(img_prompt, mask_prompt)
-            else:
-                # img_prompt, mask_prompt = [], []
-                # for key in keys:
-                #     img_prompt.append(f'a photo of {organ} image, with {kind}.')
-                #     mask_prompt.append(f'a photo of {organ} label, with {kind}.')
-                # img_prompt = img_prompt * num_samples
-                # mask_prompt = mask_prompt * num_samples
-                img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
-                mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
-
-        else:
-
-            img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
-            mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
-            # keys = prompts['key']
+        img_prompt = [f'a photo of {organ} image, with {kind}.'] * num_samples
+        mask_prompt = [f'a photo of {organ} label, with {kind}.'] * num_samples
+        # keys = prompts['key']
 
         with torch.inference_mode():
             img_prompt_embeds_, img_negative_prompt_embeds_ = self.pipe.encode_prompt(