Merge pull request #71 from Aydinhamedi/Beta-b

Beta b

Tools/RSNA_DICOM_data.py

@@ -9,13 +9,15 @@
 dcm_dir = 'stage_2_train_images'
 
 # Read the CSV file
-df = pd.read_csv('stage_2_train_labels.csv')
+df = pd.read_csv('stage_2_detailed_class_info.csv')
 
-# Directories for the two classes
+# Directories for the three classes
+not_normal_dir = 'database/Not Normal'
 normal_dir = 'database/NORMAL'
 pneumonia_dir = 'database/PNEUMONIA'
 
 # Create the directories if they don't exist
+os.makedirs(not_normal_dir, exist_ok=True)
 os.makedirs(normal_dir, exist_ok=True)
 os.makedirs(pneumonia_dir, exist_ok=True)
 
@@ -34,12 +36,15 @@
     # Convert the pixel array to an image
     img = Image.fromarray(pixels)
     # Get the label for this file
-    label = df[df['patientId'] == filename[:-4]]['Target'].values[0]
+    label = df[df['patientId'] == filename[:-4]]['class'].values[0]
     # Save the image to the appropriate directory with the original filename
-    if label == 0:
+    if label == 'No Lung Opacity / Not Normal':
+        img.save(os.path.join(not_normal_dir, filename[:-4] + '.jpeg'))
+        # print(f'Saved {filename[:-4]}.jpeg to {not_normal_dir}')
+    elif label == 'Normal':
         img.save(os.path.join(normal_dir, filename[:-4] + '.jpeg'))
         # print(f'Saved {filename[:-4]}.jpeg to {normal_dir}')
-    else:
+    else:  # 'Lung Opacity'
         img.save(os.path.join(pneumonia_dir, filename[:-4] + '.jpeg'))
         # print(f'Saved {filename[:-4]}.jpeg to {pneumonia_dir}')
     # Update the progress bar